"""
MCP客户端工具，用于连接MCP服务并处理查询
"""
import asyncio
from contextlib import AsyncExitStack
import json
import logging
from mcp import ClientSession
from mcp.client.sse import sse_client
from typing import Optional, List, Dict, Any, Union, Tuple
from openai import AsyncOpenAI
import datetime
import os
import sys

from config.settings import MODEL, OPENAI_API_KEY, BASE_URL
from utils.logger import setup_logger
from models.schemas import Message
from utils.prompts import get_system_prompt

class MCPWorker:
    """MCP工作器，用于连接MCP服务器并处理用户查询"""
    
    def __init__(self):
        """初始化MCP客户端"""
        self.session: Optional[ClientSession] = None
        self._session_context = None
        self._streams_context = None
        self.client = AsyncOpenAI(
            base_url=BASE_URL,
            api_key=OPENAI_API_KEY,
        )
        self.model = MODEL
        # 使用自动文件日志功能
        self.logger = setup_logger(name="mcp_worker", level="debug")

    async def connect_to_sse_server(self, server_url: str):
        """连接到MCP服务器"""
        self.logger.info(f"正在连接到MCP服务器: {server_url}")
        try:
            # 创建SSE连接
            self._streams_context = sse_client(url=server_url)
            streams = await self._streams_context.__aenter__()
            self.logger.debug(f"已创建SSE流连接")

            # 创建会话
            self._session_context = ClientSession(*streams)
            self.session: ClientSession = await self._session_context.__aenter__()
            self.logger.debug(f"已创建ClientSession")

            # 初始化会话
            await self.session.initialize()
            self.logger.info(f"已成功初始化MCP会话")
            
            # 列出可用工具
            response = await self.session.list_tools()
            tools = response.tools
            tool_names = [tool.name for tool in tools]
            self.logger.info(f"可用工具: {', '.join(tool_names)}")
            return tools
        except Exception as e:
            self.logger.error(f"连接MCP服务器失败: {str(e)}", exc_info=True)
            raise

    async def cleanup(self):
        """清理会话和连接"""
        self.logger.info("正在清理MCP连接")
        try:
            if self._session_context:
                await self._session_context.__aexit__(None, None, None)
                self.logger.debug("已清理会话上下文")
            if self._streams_context:
                await self._streams_context.__aexit__(None, None, None)
                self.logger.debug("已清理流上下文")
            self.logger.info("已成功清理所有MCP连接")
        except Exception as e:
            self.logger.error(f"清理MCP连接时出错: {str(e)}", exc_info=True)

    async def list_tools(self):
        """列出可用工具"""
        if not self.session:
            self.logger.warning("尝试列出工具，但未连接到MCP服务器")
            return []
        
        try:
            self.logger.debug("获取可用工具列表")
            response = await self.session.list_tools()
            tools = response.tools
            self.logger.debug(f"找到 {len(tools)} 个可用工具")
            return tools
        except Exception as e:
            self.logger.error(f"获取工具列表失败: {str(e)}", exc_info=True)
            return []

    async def process_query(self, query: str, history: List[Message] = None) -> Dict[str, Any]:
        """
        处理用户查询，使用大模型理解用户意图并调用相应工具
        
        参数:
            query: 用户查询
            history: 对话历史
            
        返回:
            处理结果
        """
        if not self.session:
            self.logger.error("尝试处理查询，但未连接到MCP服务器")
            return {"error": "未连接到MCP服务器"}
            
        self.logger.info(f"处理用户查询: {query}")
        
        # 准备消息历史
        if history is None:
            history = []
            
        # 添加用户消息
        messages = [m.dict(exclude_unset=True) for m in history if m.role != "system"]
        
        try:
            # 获取可用工具
            self.logger.debug("获取可用工具列表")
            response = await self.session.list_tools()
            tools = response.tools
            self.logger.debug(f"找到 {len(tools)} 个可用工具")
            
            # 构建动态系统提示词
            system_prompt = self._build_system_prompt(tools)
            self.logger.debug(f"生成系统提示词: {system_prompt[:100]}...")
            
            # 添加系统提示到消息历史
            messages.insert(0, {
                "role": "system",
                "content": system_prompt
            })
            
            # 直接添加用户查询，不使用关键字匹配检测时间相关查询
            # 系统提示中已经包含了完整的时间信息和时间处理指南
            # 大模型将根据提示进行智能推理
            messages.append({
                "role": "user",
                "content": query
            })

            # 构建可用工具描述
            available_tools = [{
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.inputSchema
                }
            } for tool in tools]

            # 发送初始请求给大模型
            self.logger.debug(f"发送请求给大模型 {self.model}")
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=available_tools
            )
            self.logger.debug("收到大模型响应")
            
            # 处理响应和工具调用
            tool_calls = []
            result_messages = []

            message = response.choices[0].message
            self.logger.debug(f"大模型返回消息: {message.content}")
            result_messages.append({
                "role": "assistant",
                "content": message.content or ""
            })

            # 处理工具调用，可能有多轮
            tool_call_count = 0
            while message.tool_calls:
                tool_call_count += 1
                self.logger.info(f"处理第 {tool_call_count} 轮工具调用")
                # 处理每个工具调用
                for tool_call in message.tool_calls:
                    tool_name = tool_call.function.name
                    tool_args = json.loads(tool_call.function.arguments)
                    
                    self.logger.info(f"调用工具: {tool_name}")
                    self.logger.debug(f"工具参数: {tool_args}")
                    
                    # 执行工具调用
                    try:
                        result = await self.session.call_tool(tool_name, tool_args)
                        self.logger.debug(f"工具调用结果: {result.content}")
                        tool_calls.append({
                            "tool": tool_name, 
                            "args": tool_args,
                            "result": result.content
                        })
                        
                        # 添加工具调用和结果到消息历史
                        messages.append({
                            "role": "assistant",
                            "tool_calls": [
                                {
                                    "id": tool_call.id,
                                    "type": "function",
                                    "function": {
                                        "name": tool_name,
                                        "arguments": json.dumps(tool_args)
                                    }
                                }
                            ]
                        })
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": str(result.content)
                        })
                    except Exception as e:
                        self.logger.error(f"调用工具 {tool_name} 失败: {str(e)}", exc_info=True)
                        tool_calls.append({
                            "tool": tool_name, 
                            "args": tool_args,
                            "error": str(e)
                        })
                        
                        # 添加错误信息到消息历史
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": f"Error: {str(e)}"
                        })

                # 获取下一个响应
                self.logger.debug("将工具调用结果发送给大模型")
                response = await self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    tools=available_tools
                )
                
                message = response.choices[0].message
                if message.content:
                    self.logger.debug(f"大模型回复: {message.content}")
                    result_messages.append({
                        "role": "assistant",
                        "content": message.content
                    })

            # 返回处理结果
            self.logger.info(f"查询处理完成，共进行了 {tool_call_count} 轮工具调用")
            return {
                "messages": result_messages,
                "tool_calls": tool_calls
            }
        except Exception as e:
            self.logger.error(f"处理查询时出错: {str(e)}", exc_info=True)
            return {
                "error": f"处理查询时出错: {str(e)}",
                "messages": [{"role": "assistant", "content": f"处理查询时发生错误: {str(e)}"}],
                "tool_calls": []
            }

    def _build_system_prompt(self, tools) -> str:
        """
        构建系统提示词，包含所有可用工具的说明
        
        参数:
            tools: 可用工具列表
            
        返回:
            系统提示词
        """
        return get_system_prompt(tools)